import torch
import argparse
import yaml
from geom.pc_encoder import load_geom_encoder
from dataset import read_data, GraphDataset
from model import GeoCA3D
from cfd.models.utils import MLP
import time
import numpy as np
from torch_geometric.loader import DataLoader
import os

class LpLoss(object):
    def __init__(self, d=2, p=2, size_average=True, reduction=True):
        super(LpLoss, self).__init__()

        # Dimension and Lp-norm type are postive
        assert d > 0 and p > 0

        self.d = d
        self.p = p
        self.reduction = reduction
        self.size_average = size_average

    def abs(self, x, y):
        num_examples = x.size()[0]

        # Assume uniform mesh
        h = 1.0 / (x.size()[1] - 1.0)

        all_norms = (h ** (self.d / self.p)) * torch.norm(
            x.view(num_examples, -1) - y.view(num_examples, -1), self.p, 1
        )

        if self.reduction:
            if self.size_average:
                return torch.mean(all_norms)
            else:
                return torch.sum(all_norms)

        return all_norms

    def rel(self, x, y):
        num_examples = x.size()[0]

        diff_norms = torch.norm(
            x.reshape(num_examples, -1) - y.reshape(num_examples, -1), self.p, 1
        )
        y_norms = torch.norm(y.reshape(num_examples, -1), self.p, 1)

        if self.reduction:
            if self.size_average:
                return torch.mean(diff_norms / y_norms)
            else:
                return torch.sum(diff_norms / y_norms)

        return diff_norms / y_norms

    def __call__(self, x, y):
        return self.rel(x, y)


parser = argparse.ArgumentParser()
parser.add_argument('--data_dir', default='data/Track_B')
parser.add_argument('--info_dir', default='data/Track_B')
parser.add_argument('--extra_dir', default=None)
parser.add_argument('--fold_id', default=0, type=int)
parser.add_argument('--gpu', default=0, type=int)
parser.add_argument('--val_iter', default=10, type=int)
parser.add_argument('--config_dir', default='params.yaml')
parser.add_argument('--ulip_model', default='ULIP_PointBERT')
parser.add_argument('--ulip_ckpt', default='../../GeoCA3D_init/geom/ckpt/pretrained_models_ckpt_zero-sho_classification_checkpoint_pointbert.pt')
parser.add_argument('--frozen', default=True)
parser.add_argument('--cfd_config_dir', default='cfd/cfd_params.yaml')
parser.add_argument('--cfd_model', default='MLP')
parser.add_argument('--cfd_mesh', default=True)
parser.add_argument('--weight', default=0.5, type=float)
parser.add_argument('--addTestSet', default=False)
parser.add_argument('--loadPretrainedMLP', default=False)
parser.add_argument('--PretrainedMLPPath', default=None)

args = parser.parse_args()
print(args)

with open(args.cfd_config_dir, 'r') as f:
    cfd_hparams = yaml.safe_load(f)[args.cfd_model]
print(cfd_hparams)

train_data, val_data, val_index, coef_norm = read_data(args, norm=True)

use_height = False
r = cfd_hparams['r'] if 'r' in cfd_hparams.keys() else None
train_ds = GraphDataset(train_data, use_height=use_height, use_cfd_mesh=args.cfd_mesh, r=r)
val_ds = GraphDataset(val_data, use_height=use_height, use_cfd_mesh=args.cfd_mesh, r=r)

with open(args.config_dir, 'r') as f:
    hparams = yaml.safe_load(f)['GeoCA3D']

n_gpu = torch.cuda.device_count()
use_cuda = 0 <= args.gpu < n_gpu and torch.cuda.is_available()
device = torch.device(f'cuda:{args.gpu}' if use_cuda else 'cpu')

if args.ulip_model == 'none':
    g_encoder, g_proj = None, None
else:
    g_encoder, g_proj = load_geom_encoder(args, pretrained=True, frozen=args.frozen)
    print(hparams)

encoder = MLP(cfd_hparams['encoder'], batch_norm=False)
decoder = MLP(cfd_hparams['decoder'], batch_norm=False)

if args.cfd_model == 'GraphSAGE':
    from cfd.models.GraphSAGE import GraphSAGE
    model = GraphSAGE(cfd_hparams, encoder, decoder)
elif args.cfd_model == 'MLP':
    from cfd.models.NN import NN
    model = NN(cfd_hparams, encoder, decoder)
elif args.cfd_model == 'GAT':
    from cfd.models.GAT import GAT
    model = GAT(cfd_hparams, encoder, decoder)
elif args.cfd_model == 'GNO':
    from cfd.models.GNO import GNO
    model = GNO(cfd_hparams, encoder, decoder)

model = GeoCA3D(model, geom_encoder=g_encoder, geom_proj=g_proj, **hparams)
path = 'metrics/model_0.13985563115390126.pth'
model = torch.load(path).to(device)

val_loader = DataLoader(val_ds, batch_size=1)

if not os.path.exists('output/answer'):
    os.makedirs('output/answer')
with torch.no_grad():
    model.eval()

    criterion_func = LpLoss(size_average=True)
    l2errs_press = []
    # times = []
    index = 0
    for cfd_data, geom in val_loader:
        cfd_data = cfd_data.to(device)
        geom = geom.to(device)
        out = model((cfd_data, geom))
        targets = cfd_data.y

        if coef_norm is not None:
            mean = torch.tensor(coef_norm[2]).to(device)
            std = torch.tensor(coef_norm[3]).to(device)
            pred_press = out * std[-1] + mean[-1]
            gt_press = targets * std[-1] + mean[-1]
            l2_loss = criterion_func(pred_press.unsqueeze(0), gt_press.unsqueeze(0).unsqueeze(-1))
            print('l2 loss:', l2_loss.item())
            l2errs_press.append(l2_loss.detach().cpu().numpy())
            np.save(f"./output/answer/press_{val_index[index]}.npy", pred_press.detach().cpu().numpy().squeeze())
            index = index + 1

    print('l2errs_press:', np.mean(l2errs_press))

